import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import json


def draw_chinese_text(img, text, position, font_size=24, color=(255, 255, 255), font_path="simhei.ttf"):
    img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)
    try:
        font = ImageFont.truetype(font_path, font_size)
    except:
        print(f"无法加载字体，请确认路径是否正确: {font_path}")
        font = ImageFont.load_default()
    draw.text(position, text, font=font, fill=(color[2], color[1], color[0]))
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)


class LaneCalibrator:
    def __init__(self, bird_view_size=(1920, 1020), font_path="simhei.ttf"):
        self.width, self.height = bird_view_size
        self.left_lane_points = []
        self.right_lane_points = []
        self.center_lane_points = []
        self.calibrated = False
        self.current_side = "left"
        self.calibration_frame = None
        self.font_path = font_path

    def calibrate(self, reference_image=None):
        print("开始车道线校准...")
        cv2.namedWindow("CalibrationWindow", cv2.WINDOW_NORMAL)  # 允许窗口调整大小
        cv2.resizeWindow("CalibrationWindow", 1920, 1020)  # 设置窗口大小
        cv2.setMouseCallback("CalibrationWindow", self._mouse_callback)

        self.calibration_frame = np.zeros((self.height, self.width, 3), dtype=np.uint8)

        if reference_image is not None:
            ref_h, ref_w = reference_image.shape[:2]
            if ref_w != self.width or ref_h != self.height:
                reference_image = cv2.resize(reference_image, (self.width, self.height))
            alpha = 0.5
            self.calibration_frame = cv2.addWeighted(reference_image, alpha, self.calibration_frame, 1 - alpha, 0)

        while True:
            frame = self.calibration_frame.copy()

            # 文字说明
            frame = draw_chinese_text(frame, "车道线校准工具", (150, 30), 28, (0, 200, 255), self.font_path)
            frame = draw_chinese_text(frame, "1. 标注左侧车道线 (从上到下)", (50, 70), 22, (0, 255, 0), self.font_path)
            frame = draw_chinese_text(frame, "2. 按空格键切换到右侧车道线", (50, 100), 22, (0, 255, 0), self.font_path)
            frame = draw_chinese_text(frame, "3. 标注右侧车道线 (从上到下)", (50, 130), 22, (0, 255, 0), self.font_path)
            frame = draw_chinese_text(frame, "4. 按回车键完成校准", (50, 160), 22, (0, 255, 0), self.font_path)
            frame = draw_chinese_text(frame, "5. 按ESC键取消校准", (50, 190), 22, (0, 0, 255), self.font_path)

            side_text = f"当前标注: {'左侧' if self.current_side == 'left' else '右侧'}"
            point_text = f"左侧点: {len(self.left_lane_points)}  右侧点: {len(self.right_lane_points)}"
            frame = draw_chinese_text(frame, side_text, (50, 220), 22,
                                      (0, 255, 0) if self.current_side == "left" else (0, 0, 255), self.font_path)
            frame = draw_chinese_text(frame, point_text, (50, 250), 22, (200, 200, 0), self.font_path)

            for point in self.left_lane_points:
                cv2.circle(frame, point, 5, (0, 255, 0), -1)
            for point in self.right_lane_points:
                cv2.circle(frame, point, 5, (0, 0, 255), -1)

            if len(self.left_lane_points) > 1:
                cv2.polylines(frame, [np.array(self.left_lane_points)], False, (0, 200, 0), 2)
            if len(self.right_lane_points) > 1:
                cv2.polylines(frame, [np.array(self.right_lane_points)], False, (0, 0, 200), 2)

            cv2.imshow("CalibrationWindow", frame)
            key = cv2.waitKey(1) & 0xFF

            if key == ord(' '):
                self.current_side = "right" if self.current_side == "left" else "left"
                print(f"切换到{'右侧' if self.current_side == 'right' else '左侧'}标注")
            elif key == 13:  # Enter
                if len(self.left_lane_points) >= 3 and len(self.right_lane_points) >= 3:
                    self._generate_smooth_lanes()
                    self.calibrated = True
                    print("校准完成！")
                    break
                else:
                    print("请至少标注3个左侧点和3个右侧点")
            elif key == 27:  # ESC
                print("取消校准")
                break

        cv2.destroyAllWindows()
        return self.calibrated

    def _mouse_callback(self, event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            if self.current_side == "left":
                self.left_lane_points.append((x, y))
                print(f"添加左侧点: ({x}, {y})")
            else:
                self.right_lane_points.append((x, y))
                print(f"添加右侧点: ({x}, {y})")

    def _generate_smooth_lanes(self):
        self.left_lane_points.sort(key=lambda p: p[1])
        left_x = [p[0] for p in self.left_lane_points]
        left_y = [p[1] for p in self.left_lane_points]
        left_fit = np.polyfit(left_y, left_x, 2)
        self.left_lane_points = [(int(np.polyval(left_fit, y)), y) for y in range(100, self.height - 100, 10)]

        self.right_lane_points.sort(key=lambda p: p[1])
        right_x = [p[0] for p in self.right_lane_points]
        right_y = [p[1] for p in self.right_lane_points]
        right_fit = np.polyfit(right_y, right_x, 2)
        self.right_lane_points = [(int(np.polyval(right_fit, y)), y) for y in range(100, self.height - 100, 10)]

        self.center_lane_points = []
        for l, r in zip(self.left_lane_points, self.right_lane_points):
            if l[1] == r[1]:
                cx = (l[0] + r[0]) // 2
                self.center_lane_points.append((cx, l[1]))
        print("车道线已平滑处理")

    def save_calibration(self, filepath):
        if not self.calibrated:
            print("尚未完成校准，无法保存")
            return False
        data = {
            "left_lane_points": self.left_lane_points,
            "right_lane_points": self.right_lane_points,
            "center_lane_points": self.center_lane_points
        }
        try:
            with open(filepath, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=4)
            print(f"校准结果保存到 {filepath}")
            return True
        except Exception as e:
            print(f"保存校准结果失败: {e}")
            return False
